load("//bazel:api.bzl", "mojo_filecheck_test", "mojo_test")

_EXTRA_CONSTRAINTS = {
    "test_cudnn_conv.mojo": ["//:nvidia_gpu"],
    "test_conv_transposed_cudnn.mojo": ["//:nvidia_gpu"],
    "test_irfft.mojo": ["//:nvidia_gpu"],
    "test_flash_attention_amd.mojo": ["//:amd_gpu"],
    "test_mha_causal_mask_amd.mojo": ["//:amd_gpu"],  # FIXME: KERN-1448, KERN-1437, KERN-1429
    "test_mha_operand_tma.mojo": select({
        "//:h100_gpu": [],
        "//:b200_gpu": [],
        "//conditions:default": ["@platforms//:incompatible"],
    }),
    "test_mla.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),
    "test_score_mod_mha.mojo": ["//:nvidia_gpu"],  # FIXME: KERN-1377
    "test_fused_qk_rope.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2366
    "test_layer_norm.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2397
    "test_rms_norm_fused_residual_add.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "test_group_norm.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2397
    "test_topk_gpu.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "test_rms_norm.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "test_winograd_conv.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "test_toppminp_gpu.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "test_argmaxmin_gpu.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2397
    "test_gather.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2366
    "test_gather_nd_oob.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2366
    "test_mha_tile_scheduler.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2405
    "test_fused_qk_rope_ragged.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "test_argsort.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2366
    "test_mha_mask.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: KERN-2031
    "test_mha_chunked_causal_mask.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "test_softmax.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "test_flash_attention.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "test_mha_causal_mask.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: PRDT-506
    "test_rope_ragged.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2366
    "test_pool_gpu.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2435
    "test_padding_gpu.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: GEX-2562
    "test_moe_indices.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2489
    "test_topk_gpu_fi.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2673
}

_FILECHECK_TESTS = [
    "test_mha_tile_scheduler.mojo",
    "test_padding_gpu.mojo",
    "test_gather.mojo",
    "test_gather_nd_oob.mojo",
]

[
    mojo_filecheck_test(
        name = src + ".test",
        size = "large",
        srcs = [src],
        expect_crash = src == "test_gather_nd_oob.mojo",
        tags = ["gpu"],
        target_compatible_with = ["//:has_gpu"] + _EXTRA_CONSTRAINTS.get(src, []),
        deps = [
            "//max:internal_utils",
            "//max:linalg",
            "//max:nn",
            "//max:testdata",
            "@mojo//:stdlib",
        ],
    )
    for src in _FILECHECK_TESTS
]

[
    mojo_test(
        name = src + ".test",
        size = "large",
        srcs = [src],
        copts = [],
        enable_assertions = select({
            "//:amd_gpu": False,
            "//conditions:default": src != "test_layer_norm.mojo",
        }),
        exec_properties = {
            "test.resources:gpu-memory": "2",
        },
        tags = ["gpu"],
        target_compatible_with = ["//:has_gpu"] + _EXTRA_CONSTRAINTS.get(src, []),
        deps = [
            "//max:internal_utils",
            "//max:linalg",
            "//max:nn",
            "//max:testdata",
            "@mojo//:stdlib",
        ],
    )
    for src in glob(
        ["**/*.mojo"],
        exclude = _FILECHECK_TESTS,
    )
]
